# Written by John Hoffman
# see LICENSE.txt for license information

import urllib2
import socks

from base64 import b64encode
from httplib import HTTPConnection, HTTPSConnection, IncompleteRead
from urllib import quote, unquote
from gzip import GzipFile
from StringIO import StringIO

from bencode import bdecode
from utility import exceptionArgsToString
from __init__ import user_agent

def get_auth(user, pwd):
    user_pass = '%s:%s' % (unquote(user), unquote(pwd))
    return b64encode(user_pass).strip()            
    
def get_proxy_auth():
    p = socks._defaultproxy
    if p and p[0] and p[4] != None and p[5] != None:
        return get_auth(p[4], p[5])       
    return ""
 
class urlopen(urllib2.Request):
    def __init__(self, url, encoding = "gzip", cookies = None,
                 agent = None, basicauth = None):
        if agent is None:
            agent = user_agent
        self.response = None
        urllib2.Request.__init__(self, url, headers={'User-Agent': agent})
        if encoding:
            self.add_header('Accept-Encoding', encoding)
        if cookies:
            self.add_header('Cookie', ';'.join([cookies]))
        if basicauth:
            self.add_header('Authorization', basicauth)
        if socks._defaultproxy and socks._defaultproxy[0]:
            auth = get_proxy_auth()
            if auth and socks._defaultproxy[0] >= 3:
                self.add_header('Proxy-authorization', 'Basic ' + auth)
            self.set_proxy(socks._defaultproxy[1]+":"+str(socks._defaultproxy[2]), socks.getType(socks._defaultproxy[0]))

        self._open(url)
        self.error_return = None
    
    def __nonzero__(self):
        return bool(self.response)
    
    def _open(self, url):
        try:
            self.response = urllib2.urlopen(self)
        except urllib2.URLError, e:
            raise IOError, ('http error', exceptionArgsToString(e))
        status = self.response.code
        if status != 200:
            try:
                data = self._read()
                d = bdecode(data)
                if d.has_key('failure reason'):
                    self.error_return = data
                    return
            except:
                pass
            raise IOError, ('http error', status, self.response.msg)

    def read(self):
        if self.error_return:
            return self.error_return
        return self._read()

    def _read(self):
        try:
            data = self.response.read()
        except IncompleteRead:
            return
        
        if self.response.info().get('Content-Encoding','').find('gzip') >= 0:
            try:
                compressed = StringIO(data)
                f = GzipFile(fileobj = compressed)
                data = f.read()
            except:
                raise IOError, ('http error', 'got corrupt response')
        return data

    def close(self):
        pass

    def getHeader(self, header):
        if self.response:
            return self.response.info().get(header, '')
        return ''
    
class btHTTPcon(HTTPConnection):
    """ attempt to add automatic connection timeout for HTTP connections """       
    def connect(self):
        """Connect to the host and port specified in __init__."""
        msg = "getaddrinfo returns an empty list"
        for res in socket.getaddrinfo(self.host, self.port, 0,
                                      socket.SOCK_STREAM):
            af, socktype, proto, canonname, sa = res
            try:
                self.sock = socks.socksocket(af, socktype, proto)
                self.sock.settimeout(30)
                if self.debuglevel > 0:
                    print "connect: (%s, %s)" % (self.host, self.port)
                self.sock.connect(sa)
            except socket.error, msg:
                if self.debuglevel > 0:
                    print 'connect fail:', (self.host, self.port)
                if self.sock:
                    self.sock.close()
                self.sock = None
                continue
            break
        if not self.sock:
            raise socket.error, msg

class btHTTPScon(HTTPSConnection):
    """ attempt to add automatic connection timeout for HTTPS connections """       
    def connect(self):
        "Connect to a host on a given (SSL) port."
        sock = socks.socksocket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(30)
        sock.connect((self.host, self.port))
        ssl = socket.ssl(sock, self.key_file, self.cert_file)
        self.sock = FakeSocket(sock, ssl)
